#/usr/bin/env python
import roslib; roslib.load_manifest('stack_objects')
import rospy
from geometry_msgs.msg import PointStamped
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import base
from pr2_python.world_interface import WorldInterface
from pr2_python.planning_scene_interface import get_planning_scene_interface
from pr2_python.arm_mover import ArmMover
from pr2_python.controller_manager_client import ControllerManagerClient
from pr2_tasks.pickup_tableware import PickupTableware
import copy
import ee_cart_imped_action
from pr2_python import geometry_tools
from geometry_msgs.msg import PoseStamped
from pr2_python import gripper
from pr2_python.arm_mover import ArmMover
from Room3 import *
from Placer import *

class PlatePicker:
  def __init__(self):
    wi = WorldInterface()
    global res
    global ordered_objects

  def detect_objects(self):
     global res
     det = TablewareDetection()
     res = det.detect_objects()
     rospy.loginfo('Detected objects with the following labels and poses:')
     return

  def object_diameter(self, pg):
     y_max = pg.target.cluster.points[0].y
     y_min = pg.target.cluster.points[0].y
     for point in pg.target.cluster.points:
        if point.y > y_max:
           y_max = point.y
        if point.y < y_min:
           y_min = point.y
     return (y_max-y_min)

  def separate_graspable_objects(self):
     global res      
     for pg in res.pickup_goals:
        print pg.label
        if pg.label is not 'graspable':
           pass
        else:
           diam = self.object_diameter(pg)
           print diam
           if (diam < 0.07):
              pg.label = 'small_cups'
           elif (diam < 0.14):
              pg.label = 'cups'
           elif (diam < 0.2):
              pg.label = 'bowls'
           elif (diam < 0.35):
              pg.label = 'plates'
           else:
              pg.label = 'tablet'   

  def create_stack_list(self):
     global res
     global ordered_objects
     ordered_objects = []
     known_objects = ['plate', 'bowl', 'cup']
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18980: #big sushi plate
              ordered_objects.append(pg)
        if pg.label == 'plates':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18979: #little sushi plate
              ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18976: #big sushi bowl
              ordered_objects.append(pg)
        if pg.label == 'bowls':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18975: #little sushi bowl
              ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18978: #big sushi cup
              ordered_objects.append(pg)
        if pg.label == 'cups':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18977: #little sushi cup
              ordered_objects.append(pg)
        if pg.label == 'small_cups':
           ordered_objects.append(pg)
     for pg in ordered_objects:
        rospy.loginfo('Label: '+str(pg.label))  

  def pick_top_plate(self, tasks):
     global ordered_objects
     pi = 3.14159265
     for pg in ordered_objects:
        if pg.label is 'plates':
           print "trying to pick up a plate from stack of plates"
           arm_mover_r = ArmMover('right_arm')
           arm_mover_l = ArmMover('left_arm')
           gr_r = gripper.Gripper('right_arm')
           gr_l = gripper.Gripper('left_arm')
           gr_r.open()
           gr_l.open() 

           # get the diameter of the object
           pg_max_dist= self.object_diameter(pg)

           # preset the right arm
           goal_pose = PoseStamped()
           goal_pose.pose.position.x = pg.object_pose_stamped.pose.position.x
           goal_pose.pose.position.y = pg.object_pose_stamped.pose.position.y
           goal_pose.pose.position.z = pg.object_pose_stamped.pose.position.z + 0.3
           goal_pose.pose.orientation = geometry_tools.euler_to_quaternion(0, pi/2, pi/2)
           print 'pg object frame is ' +str(pg.object_pose_stamped.header.frame_id)
           goal_pose.header.frame_id = '/base_link'

           reached_goal = arm_mover_r.move_to_goal(goal_pose)
           if not reached_goal:
              print  arm_mover_r.get_exceptions()

           # preset the left arm
           goal_pose.pose.position.x = pg.object_pose_stamped.pose.position.x
           goal_pose.pose.position.y = pg.object_pose_stamped.pose.position.y + 0.5
           goal_pose.pose.position.z = pg.object_pose_stamped.pose.position.z + 0.2
           goal_pose.pose.orientation.x = 0.5
           goal_pose.pose.orientation.y = -0.5
           goal_pose.pose.orientation.z = -0.5
           goal_pose.pose.orientation.w = 0.5

           reached_goal = arm_mover_l.move_to_goal(goal_pose)
           if not reached_goal:
              print  arm_mover_l.get_exceptions()

           # start the controllers
           control_manager = ControllerManagerClient()
           control_manager.switch_controllers(start_controllers=['r_arm_cart_imped_controller','l_arm_cart_imped_controller'],stop_controllers=['r_arm_controller','l_arm_controller'])

           r_control = ee_cart_imped_action.EECartImpedClient('right_arm')
           l_control = ee_cart_imped_action.EECartImpedClient('left_arm')
           # slide top plate way
           r_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                      pg.object_pose_stamped.pose.position.y,
                                      pg.object_pose_stamped.pose.position.z + 0.1,
                                      -0.5, 0.5, 0.5, 0.5,
                                      1000, 1000, 1000, 30, 30, 30,
                                      False, False, False, False, False,
                                      False, 1, 
                                      '/base_link');
           r_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                      pg.object_pose_stamped.pose.position.y,
                                      pg.object_pose_stamped.pose.position.z,
                                      -0.5, 0.5, 0.5, 0.5,
                                      1000, 1000, -2, 30, 30, 30,
                                      False, False, True, False, False,
                                      False, 3, 
                                      '/base_link');
           r_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                      pg.object_pose_stamped.pose.position.y+0.07,
                                      pg.object_pose_stamped.pose.position.z,
                                      -0.5, 0.5, 0.5, 0.5,
                                      1000, 1000, -5, 30, 30, 30,
                                      False, False, True, False, False,
                                      False, 6, 
                                      '/base_link');
           r_control.sendGoal()

           # move l arm closer
           l_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                        pg.object_pose_stamped.pose.position.y + 0.25,
                                        pg.object_pose_stamped.pose.position.z + 0.05,
                                      0.5, -0.5, -0.5, 0.5,
                                      1000, 1000, 1000, 30, 30, 30,
                                      False, False, False, False, False,
                                      False, 4, 
                                      '/base_link');
           l_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                        pg.object_pose_stamped.pose.position.y + ((pg_max_dist/2.0)-.01),
                                        pg.object_pose_stamped.pose.position.z + 0.02,
                                      0.5, -0.5, -0.5, 0.5,
                                      1000, -3, 1000, 30, 30, 30,
                                      False, True, False, False, False,
                                      False, 6, 
                                      '/base_link');
           l_control.sendGoal()
           gr_l.close()
           l_control.cancelAllGoals()

           # move right arm out of collision        
           r_control.resetGoal() 
           r_control.addTrajectoryPoint(pg.object_pose_stamped.pose.position.x, 
                                      pg.object_pose_stamped.pose.position.y,
                                      pg.object_pose_stamped.pose.position.z+0.3,
                                      -0.5, 0.5, 0.5, 0.5,
                                      1000, 1000, 1000, 30, 30, 30,
                                      False, False, False, False, False,
                                      False, 2, 
                                      '/base_link');
  
           r_control.sendGoal()
  
           # move back to the old controllers
           control_manager.switch_controllers(stop_controllers=['r_arm_cart_imped_controller','l_arm_cart_imped_controller'],start_controllers=['r_arm_controller','l_arm_controller'])

           # add the object to the world in the gripper 
           world =  WorldInterface()
           world.add_collision_box(pg.object_pose_stamped, (0.26, 0.26, 0.02), 'plate')
           world.attach_object_to_gripper('right_arm','plate') 
           # and get ready for next action
           tasks.move_arm_to_side('right_arm')

  def move_arms_back(self):
    tasks = Tasks() 
    tasks.move_arms_to_side()

  def pick_plate_stack(self):
    global ordered_objects
    self.reset()
    rospy.sleep(1.0)

    self.move_arms_back()

    gr_r = gripper.Gripper('right_arm')
    gr_l = gripper.Gripper('left_arm')
    gr_r.open()
    gr_l.open()

    self.detect_objects()
    self.separate_graspable_objects()
    self.create_stack_list()

    pt = PickupTableware()
    for pg in ordered_objects:
       if pg.label is 'plates' or pg.label is 'plate':
          print "trying to pick up a plate"
          pt.grasp_plate(pg)

          tasks = Tasks()
          tasks.move_arm_to_side('left_arm')

          arm_mover_r = ArmMover('right_arm')

          # preset the right arm
          goal_pose = PoseStamped()
          goal_pose.pose.position.x = pg.object_pose_stamped.pose.position.x
          goal_pose.pose.position.y = pg.object_pose_stamped.pose.position.y
          goal_pose.pose.position.z = pg.object_pose_stamped.pose.position.z + 0.3
          goal_pose.pose.orientation = geometry_tools.euler_to_quaternion(-3.14/2, 0.0, 0.0)
          print 'pg object frame is ' +str(pg.object_pose_stamped.header.frame_id)
          goal_pose.header.frame_id = '/base_link'

          reached_goal = arm_mover_r.move_to_goal(goal_pose)
          if not reached_goal:
             print  arm_mover_r.get_exceptions()


  def put_on_table(self):
    tasks = Tasks() 
    place_stamped = PoseStamped()
    place_stamped.header.frame_id = '/map'
    place_stamped.pose.position.x = -1.1
    place_stamped.pose.position.y = 0.2
    place_stamped.pose.position.z = 0.75
    place_stamped.pose.orientation.w = 1.0

    rospy.loginfo('Going to place!')
    tasks.go_and_place('left_arm', place_stamped)

  def drop_stack_on_table(self):
    placer = Placer()
    tasks = Tasks()
    place_stamped = PoseStamped()
    place_stamped.header.frame_id = '/map'
    place_stamped.pose.position.x = -1.4
    place_stamped.pose.position.y = 0.2
    place_stamped.pose.position.z = 0.85
    #place_stamped.pose.orientation = geometry_tools.euler_to_quaternion(-3.14/2, 0.0, 0.0)
    place_stamped.pose.orientation = geometry_tools.euler_to_quaternion(0.0, 0.0, 0.0)

    rospy.loginfo('Going to place!')
    placer.set_it_there('right_arm', place_stamped)
    #tasks.go_and_place('right_arm', place_stamped)

  def reset(self):
    # reset the world
    wi = WorldInterface()
    wi.reset()
    psi = get_planning_scene_interface()
    psi.reset()

  def put_plate_from_stack_on_table(self):
    tasks = Tasks() 
    self.reset()
    self.move_arms_back()
    #room = Room3()
    #room.goto_table()
    self.detect_objects()
    self.separate_graspable_objects()
    self.create_stack_list()

    psi = get_planning_scene_interface()
    psi.reset()
    rospy.sleep(0.5)
    print "Press a key when you're ready to start picking plates"
    raw_input()

    self.pick_top_plate(tasks)
    # TODO UNCOMMENT AGAIN WHEN MAP IS READY
    #self.put_on_table()
